import cv2
import numpy as np
class perspectiveTransform:
    def __init__(self,x,y,debug = False):
        self.x = x
        self.y = y
        self.org_x  = 0
        self.org_y = 0
        self.debug = debug
        self.calib_matrix = None
        self.chessboard_corners = None
        
    def calibration(self, calib_img_chessboad,calib_img_frame,chessboardSize):
        if self.chessboard_corners == None:
            chessboard_loaded_correctly, chessboard_corners = self._get_Corners(calib_img_chessboad,chessboardSize)
            if chessboard_loaded_correctly:
                self.org_y = calib_img_chessboad.shape[:2][0]
                self.org_x = calib_img_chessboad.shape[:2][1]
                self.chessboard_corners = chessboard_corners
        frame_loaded_correctly,chessboard_frame = self._get_Corners(calib_img_frame,chessboardSize)
        if frame_loaded_correctly:
            if self.chessboard_corners != None:
                self.calib_matrix = cv2.getPerspectiveTransform(chessboard_frame,self.chessboard_corners)
                return True
        return False
        
    def rectifie_frame(self,frame):
        if(self.calib_matrix != None):
            img = cv2.warpPerspective(frame,self.calib_matrix,(self.org_x,self.org_y))
            return cv2.resize(img, (self.x,self.y))
        else:
            raise Exception("No calib matrix loaded")
    
    def _get_Corners(self, frame, chessboardSize):
        founded, corners = cv2.findChessboardCorners(frame, chessboardSize)
        if(founded):
            rh = (corners[0][0][0],corners[0][0][1])
            lh = (corners[len(corners) -chessboardSize[0]][0][0],corners[len(corners) -chessboardSize[0]][0][1]) 
            rl = (corners[chessboardSize[0] -1][0][0],corners[chessboardSize[0] -1][0][1])
            ll = (corners[len(corners) -1][0][0],corners[len(corners) -1][0][1]) 
            return_value =np.array([rh,lh,rl,ll])
            if self.debug and founded:            
                cv2.drawChessboardCorners(frame, chessboardSize, corners, founded)
                cv2.circle(frame, (int(rh[0]),int(rh[1])), 2, (255,0,0),3)
                cv2.circle(frame, (int(lh[0]),int(lh[1])), 2, (0,255,0),3)
                cv2.circle(frame, (int(rl[0]),int(rl[1])), 2, (0,0,255),3)
                cv2.circle(frame, (int(ll[0]),int(ll[1])), 2, (255,255,255),3)
                cv2.imshow("chessboard",frame)
                cv2.waitKey(0)
                cv2.destroyWindow("chessboard")
                
            return True, return_value
        else:
            return False, None
        
        